import os
import pandas as pd
import re
from openpyxl import load_workbook
from openpyxl.styles import PatternFill
import logging
from multiprocessing import Pool, cpu_count

# Set up logging
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

class DataCleaner:
    """
    数据清洗类，用于批量处理Excel文件，筛选无效数据并标红被删除的行。
    """

    def __init__(self, raw_data_folder, clean_data_folder):
        """
        初始化DataCleaner对象。
        :param raw_data_folder: 原始数据文件夹路径。
        :param clean_data_folder: 清洗后数据文件夹路径。
        """
        self.raw_data_folder = raw_data_folder
        self.clean_data_folder = clean_data_folder
        os.makedirs(clean_data_folder, exist_ok=True)

        # 定义噪声文本的规则
        self.noise_patterns = [
            r"^(好|行|你|再见|拜拜|早上好|晚上好)[，。；？！]*$",  # 单个字后面跟标点符号
            r"^\d+[，。；？！]*$",  # 单独的阿拉伯数字后面跟标点符号
            r"^(嗯+|啊+|哦+|哈+)[，。；？！]*$",  # 语气词
            r"^\s*$"  # 空白行或空字符串
        ]

    def is_invalid_text(self, text):
        """
        判断文本是否为无效内容。
        :param text: 文本内容字符串。
        :return: 如果是无效内容返回True，否则返回False。
        """
        for pattern in self.noise_patterns:
            if re.match(pattern, text):
                return True
        return False

    def process_file(self, file_path):
        """
        处理单个Excel文件，筛选无效数据并保存结果。
        :param file_path: 原始文件的路径。
        """
        try:
            # 读取Excel文件
            df = pd.read_excel(file_path, header=None, names=["时间", "文本内容", "场景", "标签"])
            df["status"] = df["文本内容"].apply(lambda text: "删除" if self.is_invalid_text(text) else "保留")

            # 处理相邻的相同文本
            for index in range(1, len(df)):
                if df.loc[index, "文本内容"] == df.loc[index - 1, "文本内容"]:
                    df.at[index, "status"] = "删除"

            # 将“文本内容”和“标签”移动到前面
            cleaned_df = df[["文本内容", "标签", "时间", "场景", "status"]].rename(columns={"文本内容": "text", "标签": "label"})

            # 保存清洗后的数据
            filename = os.path.basename(file_path)
            output_path = os.path.join(self.clean_data_folder, f"cleaned_{filename}")
            cleaned_df.to_excel(output_path, index=False)

            # 标红被删除的行
            self.mark_deleted_rows(output_path, cleaned_df)
            logging.info(f"Processed file: {filename}")

        except Exception as e:
            logging.error(f"Failed to process file {file_path}: {e}")

    @staticmethod
    def mark_deleted_rows(file_path, cleaned_df):
        """
        标红被删除的行。
        :param file_path: 清洗后文件的路径。
        :param cleaned_df: 清洗后的数据框。
        """
        # 加载Excel文件
        workbook = load_workbook(file_path)
        sheet = workbook.active

        # 设置标红格式
        red_fill = PatternFill(start_color="FF0000", end_color="FF0000", fill_type="solid")
        for row_idx, status in enumerate(cleaned_df["status"], start=2):  # 跳过表头
            if status == "删除":
                for col_idx in range(1, sheet.max_column + 1):
                    sheet.cell(row=row_idx, column=col_idx).fill = red_fill

        # 保存文件
        workbook.save(file_path)

    def process_all_files(self):
        """
        批量处理raw_data_folder中的所有Excel文件。
        """
        file_paths = [os.path.join(self.raw_data_folder, filename) for filename in os.listdir(self.raw_data_folder) if filename.endswith('.xlsx')]

        with Pool(processes=cpu_count()) as pool:
            pool.map(self.process_file, file_paths)

        logging.info("所有数据清洗完成，结果已保存到 clean_data 文件夹中。")


# 主程序入口
if __name__ == "__main__":
    # 定义文件夹路径
    RAW_DATA_FOLDER = 'raw_data/'  # 原始数据文件夹
    CLEAN_DATA_FOLDER = 'clean_data/'  # 清洗后数据文件夹

    # 创建DataCleaner对象并处理所有文件
    cleaner = DataCleaner(RAW_DATA_FOLDER, CLEAN_DATA_FOLDER)
    cleaner.process_all_files()
